package cn.micro.dubbo.trace.filter;

import cn.micro.core.context.CurrentContext;
import cn.micro.core.util.TraceIdUtil;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.*;
import org.slf4j.MDC;
import org.springframework.core.Ordered;
import org.springframework.util.StringUtils;

import java.util.concurrent.atomic.AtomicInteger;

import static cn.micro.core.constant.Cants.CURRENT_TRACE_ID;
import static cn.micro.core.constant.Cants.X_CURRENT_REQUEST_ID;

@Activate(group = "provider", order = Ordered.HIGHEST_PRECEDENCE)
public class TraceProviderFilter implements Filter {

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        RpcContext rpcContext = RpcContext.getContext();
        String traceId = rpcContext.getAttachment(X_CURRENT_REQUEST_ID);
        String currentTraceId = CurrentContext.get(CURRENT_TRACE_ID, TraceIdUtil.genTraceId());
        // 如果是 dubbo 本地调用，啥也不干
        if (!StringUtils.isEmpty(traceId) && traceId.startsWith(currentTraceId)) {
            return invoker.invoke(invocation);
        }
        try {
            if (StringUtils.isEmpty(traceId)) {
                traceId = TraceIdUtil.genTraceId();
            }
            // get traceId from dubbo consumer，and set traceId to MDC
            MDC.put(X_CURRENT_REQUEST_ID, "[" + traceId + "]");
            CurrentContext.set(CURRENT_TRACE_ID, traceId);
            CurrentContext.set(traceId, new AtomicInteger(0));
            return invoker.invoke(invocation);
        } finally {
            // clear traceId from MDC
            MDC.remove(X_CURRENT_REQUEST_ID);
            CurrentContext.removeKey(traceId);
            CurrentContext.removeKey(CURRENT_TRACE_ID);
        }
    }
}
